# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import statsmodels.api as sm
from scipy.linalg import eig, svd
from statsmodels.discrete.discrete_model import NegativeBinomial
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor
import chardet
import re
import warnings
from scipy import stats

# 过滤警告信息
warnings.filterwarnings('ignore')

# 设置Seaborn的样式
sns.set(style="whitegrid")


# 定义检测文件编码的函数
def detect_encoding(file_path):
    with open(file_path, 'rb') as f:
        result = chardet.detect(f.read(100000))  # 读取前100KB
    print(f"检测到的文件编码 {file_path}: {result['encoding']}")
    return result['encoding']


# 定义文件路径（请根据实际路径进行调整）
medal_counts_path = r'data\summerOly_medal_counts_cleaned.csv'
hosts_path = r'2025_Problem_C_Data\summerOly_hosts.csv'
programs_path = r'2025_Problem_C_Data\summerOly_programs.csv'
athletes_path = r'2025_Problem_C_Data\summerOly_athletes.csv'
data_dictionary_path = r'2025_Problem_C_Data\data_dictionary.csv'


# 检测并读取CSV文件
def read_csv_with_encoding(file_path):
    try:
        encoding = detect_encoding(file_path)
        df = pd.read_csv(file_path, encoding=encoding)
    except UnicodeDecodeError:
        print(f"读取 {file_path} 时遇到 UnicodeDecodeError，尝试使用 'latin1' 编码。")
        df = pd.read_csv(file_path, encoding='latin1')
    return df


# 读取所有数据
medal_counts = read_csv_with_encoding(medal_counts_path)
hosts = read_csv_with_encoding(hosts_path)
programs = read_csv_with_encoding(programs_path)
athletes = read_csv_with_encoding(athletes_path)
data_dictionary = read_csv_with_encoding(data_dictionary_path)

# 显示各数据集的前几行
print("\n奖牌统计数据 (Medal Counts):")
print(medal_counts.head())

print("\n主办城市数据 (Hosts):")
print(hosts.head())

print("\n项目数据 (Programs):")
print(programs.head())

print("\n运动员数据 (Athletes):")
print(athletes.head())

# -------------------------------
# 2. 数据预处理
# -------------------------------

# 2.1. 清理 summerOly_programs.csv
# 替换特殊符号为 NaN 并清理注释
programs.replace({'': np.nan, '': np.nan}, inplace=True)
programs['Sports Governing Body'] = programs['Sports Governing Body'].str.replace(r'\[.*?\]', '',
                                                                                  regex=True).str.strip()

# 将项目数据从宽格式转换为长格式
years = [str(year) for year in range(1896, 2033, 4)]  # 包括2028和2032
id_vars = ['Sport', 'Discipline', 'Code', 'Sports Governing Body']
existing_years = [year for year in years if year in programs.columns]
missing_years = set(years) - set(programs.columns)
if missing_years:
    print(f"\n警告: 以下年份在项目数据中缺失，将从 melt 中排除: {sorted(missing_years)}")

value_vars = existing_years
programs_long = programs.melt(id_vars=id_vars, value_vars=value_vars, var_name='Year', value_name='Event_Count')

# 将 Year 转换为整数并处理缺失值
programs_long['Year'] = pd.to_numeric(programs_long['Year'], errors='coerce')
programs_long['Event_Count'] = programs_long['Event_Count'].astype(str).str.strip()

# 使用正则表达式移除 'Event_Count' 中的非数字字符
programs_long['Event_Count'] = programs_long['Event_Count'].apply(
    lambda x: re.findall(r'\d+', x)[0] if re.findall(r'\d+', x) else '0')

# 将 'Event_Count' 转换为整数
programs_long['Event_Count'] = programs_long['Event_Count'].fillna('0').astype(int)

# 确认 'Sport' 列存在且清理
print("\nprograms_long['Sport'] 样本:")
print(programs_long['Sport'].head())

# 去除 'Sport' 列的前后空格
programs_long['Sport'] = programs_long['Sport'].str.strip()

# 2.2. 合并 summerOly_medal_counts.csv 与 summerOly_hosts.csv
# 识别取消的奥运年份
# 通常取消的年份会在 'Host' 列中包含 'Cancelled' 字样
cancelled_years = hosts[hosts['Host'].str.contains('Cancelled', na=False)]['Year'].tolist()
print(f"\n取消的奥运年份 (Cancelled Olympic Years): {cancelled_years}")

# 仅保留非取消的年份
medal_counts = medal_counts[~medal_counts['Year'].isin(cancelled_years)]

# 创建一个国家名称到 NOC 代码的映射（根据实际数据补充）
country_to_noc = {
    "United States": "USA",
    "Greece": "GRE",
    "Germany": "DEU",
    "France": "FRA",
    "Great Britain": "GBR",  # 映射为 United Kingdom
    "China": "CHN",
    "Denmark": "DNK",
    "Netherlands": "NED",
    "Hungary": "HUN",
    "Austria": "AUT",
    "Australia": "AUS",
    "Sweden": "SWE",
    "Italy": "ITA",
    "Belgium": "BEL",
    "Switzerland": "CHE",
    "Canada": "CAN",
    "Norway": "NOR",
    "Russia": "RUS",
    "Japan": "JPN",
    "Finland": "FIN",
    "Poland": "POL",
    "Turkey": "TUR",
    "Brazil": "BRA",
    "Romania": "ROU",
    "Czechoslovakia": "CZE",  # 已映射为 CZE
    "Czech Republic": "CZE",
    "South Korea": "KOR",
    "Spain": "ESP",
    "Argentina": "ARG",
    "New Zealand": "NZL",
    "South Africa": "RSA",
    "Mexico": "MEX",
    "Iran": "IRN",
    "Ethiopia": "ETH",
    "Kenya": "KEN",
    "Ukraine": "UKR",
    "Cuba": "CUB",
    "Portugal": "PRT",
    "Azerbaijan": "AZE",
    "Bahrain": "BAH",
    "Belarus": "BLR",
    "Bulgaria": "BUL",
    "Chile": "CHL",
    "Colombia": "COL",
    "Dominican Republic": "DOM",
    "Ecuador": "ECU",
    "Egypt": "EGY",
    "Estonia": "EST",
    "Georgia": "GEO",
    "Grenada": "GRD",
    "Guatemala": "GTM",
    "Hong Kong": "HKG",
    "India": "IND",
    "Indonesia": "IDN",
    "Israel": "ISR",
    "Jamaica": "JAM",
    "Jordan": "JOR",
    "Kazakhstan": "KAZ",
    "Kosovo": "KOS",
    "Kyrgyzstan": "KGZ",
    "Latvia": "LVA",
    "Lithuania": "LTU",
    "Moldova": "MDA",
    "Mongolia": "MNG",
    "Morocco": "MAR",
    "Namibia": "NAM",
    "North Korea": "PRK",
    "North Macedonia": "MKD",  # 映射为 North Macedonia
    "Norway": "NOR",
    "Pakistan": "PAK",
    "Panama": "PAN",
    "Paraguay": "PRY",
    "Peru": "PER",
    "Philippines": "PHL",
    "Qatar": "QAT",
    "Refugee Olympic Team": "ROT",
    "San Marino": "SMR",
    "Saudi Arabia": "KSA",
    "Serbia": "SRB",  # 映射为 Serbia
    "Singapore": "SGP",
    "Slovakia": "SVK",
    "Slovenia": "SVN",
    "Syria": "SYR",
    "Tajikistan": "TJK",
    "Thailand": "THA",
    "Trinidad and Tobago": "TTO",
    "Turkmenistan": "TKM",
    "Tuvalu": "TUV",
    "Uganda": "UGA",
    "United Arab Emirates": "UAE",
    "Uruguay": "URY",
    "Uzbekistan": "UZB",
    "Vietnam": "VNM",
    "Yemen": "YEM",
    "Zambia": "ZMB",
    "Zimbabwe": "ZWE",
    "Antigua and Barbuda": "ATG",
    "Armenia": "ARM",
    "Bahamas": "BHS",
    "Bangladesh": "BGD",
    "Barbados": "BRB",
    "Benin": "BEN",
    "Bermuda": "BMU",
    "Bhutan": "BTN",
    "Bolivia": "BOL",
    "Bosnia and Herzegovina": "BIH",
    "Botswana": "BWA",
    "Brunei": "BRN",
    "Burkina Faso": "BFA",
    "Burundi": "BDI",
    "Cambodia": "KHM",
    "Cameroon": "CMR",
    "Cape Verde": "CPV",
    "Central African Republic": "CAF",
    "Chad": "TCD",
    "Comoros": "COM",
    "Congo (Brazzaville)": "COG",
    "Congo (Kinshasa)": "COD",
    "Costa Rica": "CRI",
    "Côte d'Ivoire": "CIV",
    "Croatia": "HRV",
    "Cyprus": "CYP",
    "Djibouti": "DJI",
    "Dominica": "DMA",
    "El Salvador": "SLV",
    "Equatorial Guinea": "GNQ",
    "Eritrea": "ERI",
    "Eswatini": "SWZ",  # 映射为 Swaziland
    "Fiji": "FJI",
    "Gabon": "GAB",
    "Gambia": "GMB",
    "Ghana": "GHA",
    "Grenada": "GRD",
    "Guinea": "GIN",
    "Guinea-Bissau": "GNB",
    "Guyana": "GUY",
    "Haiti": "HTI",
    "Honduras": "HND",
    "Iceland": "ISL",
    "Ireland": "IRL",
    "Kazakhstan": "KAZ",
    "Kenya": "KEN",
    "Kiribati": "KIR",
    "Kuwait": "KWT",
    "Kyrgyzstan": "KGZ",
    "Laos": "LAO",
    "Lebanon": "LBN",
    "Lesotho": "LSO",
    "Liberia": "LBR",
    "Libya": "LBY",
    "Liechtenstein": "LIE",
    "Luxembourg": "LUX",
    "Madagascar": "MDG",
    "Malawi": "MWI",
    "Malaysia": "MYS",
    "Maldives": "MDV",
    "Mali": "MLI",
    "Malta": "MLT",
    "Marshall Islands": "MHL",
    "Mauritania": "MRT",
    "Mauritius": "MUS",
    "Micronesia": "FSM",
    "Monaco": "MCO",
    "Montenegro": "MNE",
    "Mozambique": "MOZ",
    "Myanmar": "MMR",
    "Nauru": "NRU",
    "Nepal": "NPL",
    "Nicaragua": "NIC",
    "Niger": "NER",
    "Nigeria": "NGA",
    "Oman": "OMN",
    "Palau": "PLW",
    "Papua New Guinea": "PNG",
    "Qatar": "QAT",
    "Romania": "ROU",
    "Russia": "RUS",
    "Rwanda": "RWA",
    "Saint Kitts and Nevis": "KNA",
    "Saint Lucia": "LCA",
    "Saint Vincent and the Grenadines": "VCT",
    "Samoa": "WSM",
    "San Marino": "SMR",
    "Sao Tome and Principe": "STP",
    "Saudi Arabia": "KSA",
    "Senegal": "SEN",
    "Serbia": "SRB",  # 映射为 Serbia
    "Seychelles": "SYC",
    "Sierra Leone": "SLE",
    "Singapore": "SGP",
    "Slovakia": "SVK",
    "Slovenia": "SVN",
    "Solomon Islands": "SLB",
    "Somalia": "SOM",
    "South Africa": "ZAF",
    "South Korea": "KOR",
    "South Sudan": "SSD",
    "Sri Lanka": "LKA",
    "Sudan": "SDN",
    "Suriname": "SUR",
    "Sweden": "SWE",
    "Swaziland": "SWZ",  # 映射为 Eswatini
    "Switzerland": "CHE",
    "Syria": "SYR",
    "Taiwan": "TPE",  # 与noc_mapping一致
    "Tajikistan": "TJK",
    "Tanzania": "TZA",
    "Thailand": "THA",
    "Timor-Leste": "TLS",
    "Togo": "TGO",
    "Tonga": "TON",
    "Tunisia": "TUN",
    "Turkey": "TUR",
    "Turkmenistan": "TKM",
    "Tuvalu": "TUV",
    "Uganda": "UGA",
    "United Kingdom": "GBR",
    "United States": "USA",
    "Uruguay": "URY",
    "Uzbekistan": "UZB",
    "Vanuatu": "VUT",
    "Venezuela": "VEN",
    "Vietnam": "VNM",
    "Yemen": "YEM",
    "Zambia": "ZMB",
    "Zimbabwe": "ZWE"
}


# 函数：从 Host 列提取国家名称，并映射到 NOC 代码
def extract_noc(host_entry):
    if pd.isna(host_entry):
        return 'UNK'
    parts = host_entry.split(',')
    if len(parts) < 2:
        return 'UNK'
    country = parts[1].strip().replace('\xa0', ' ')
    return country_to_noc.get(country, 'UNK')


# 应用函数提取 Host_NOC
hosts['Host_NOC'] = hosts['Host'].apply(extract_noc)

# 合并 Host_NOC 到 medal_counts
medal_counts = pd.merge(medal_counts, hosts[['Year', 'Host_NOC']], on='Year', how='left')

# 创建 Host_Flag：如果 NOC 是 Host_NOC，则为1，否则为0
medal_counts['Host_Flag'] = (medal_counts['NOC'] == medal_counts['Host_NOC']).astype(int)

programs_long = programs_long.drop(columns=["Sports Governing Body", "Discipline"])


medal_counts = medal_counts.drop(columns=["Silver", "Bronze", "Rank"])

# 筛选出主办国数据
host_gold = medal_counts[medal_counts['NOC'] == medal_counts['Host_NOC']]

# 合并主办国金牌数和赛事数量
host_events = programs_long[programs_long['Year'] == programs_long['Year']].groupby(['Sport', 'Year'])['Event_Count'].sum().reset_index()

# 合并数据集：medal_counts 和 programs_long 通过 'Year' 和 'Host_NOC'
host_data = pd.merge(host_gold, host_events, on=['Year'], how='inner')

# 筛选出每个主办国的金牌数和赛事数量
host_data = pd.merge(host_gold, host_events, on=['Year'], how='inner')

# 筛选出每个主办国的金牌数和赛事数量
host_data = pd.merge(host_gold, host_events, on=['Year'], how='inner')

# 选择金牌数量前五的主办国
top_host_countries = host_data.groupby('Host_NOC')['Gold'].sum().nlargest(5).index

# 筛选出每个主办国的金牌数和赛事数量
host_data = pd.merge(host_gold, host_events, on=['Year'], how='inner')

# 选择金牌数量前五的主办国
top_host_countries = host_data.groupby('Host_NOC')['Gold'].sum().nlargest(5).index

# 筛选出前五个主办国的数据
top_host_data = host_data[host_data['Host_NOC'].isin(top_host_countries)]

# 使用pivot_table来避免重复条目问题，并聚合金牌数量
heatmap_data = top_host_data.pivot_table(index='Host_NOC', columns='Year', values='Gold', aggfunc='sum')

# 绘制散点图
plt.figure(figsize=(12, 6))
for country in top_host_countries:
    country_data = top_host_data[top_host_data['Host_NOC'] == country]
    plt.scatter(country_data['Event_Count'], country_data['Gold'], label=country)

# 设置标题和标签
plt.title("Gold Medals vs Number of Events for Top 5 Host Countries", fontsize=14)
plt.xlabel("Number of Events", fontsize=12)
plt.ylabel("Gold Medals", fontsize=12)
plt.legend(title="Host Countries", fontsize=10)

# 显示图表
plt.tight_layout()
plt.show()